package com.agentsflex.llm.deepseek;

import com.agentsflex.core.message.AiMessage;
import com.agentsflex.core.message.FunctionCall;
import com.agentsflex.core.message.MessageStatus;
import com.agentsflex.core.parser.AiMessageParser;
import com.agentsflex.core.parser.JSONObjectParser;
import com.agentsflex.core.parser.impl.DefaultAiMessageParser;
import com.agentsflex.core.util.StringUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.JSONPath;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.print.DocFlavor;
import java.util.*;

public class DeepseekAiMessageParser  implements AiMessageParser {

    private static final Logger log = LoggerFactory.getLogger(DeepseekAiMessageParser.class);
    private String contentPath;
    private String reasoningContentPath;
    private String indexPath;
    private String totalTokensPath;
    private String promptTokensPath;
    private String completionTokensPath;
    private JSONObjectParser<MessageStatus> statusParser;
    private JSONObjectParser<List<FunctionCall>> callsParser;
    private List<FunctionCall> calls;

    public String getContentPath() {
        return contentPath;
    }

    public void setContentPath(String contentPath) {
        this.contentPath = contentPath;
    }

    public String getReasoningContentPath() {
        return reasoningContentPath;
    }

    public void setReasoningContentPath(String reasoningContentPath) {
        this.reasoningContentPath = reasoningContentPath;
    }

    public String getIndexPath() {
        return indexPath;
    }

    public void setIndexPath(String indexPath) {
        this.indexPath = indexPath;
    }

    public String getTotalTokensPath() {
        return totalTokensPath;
    }

    public void setTotalTokensPath(String totalTokensPath) {
        this.totalTokensPath = totalTokensPath;
    }

    public String getPromptTokensPath() {
        return promptTokensPath;
    }

    public void setPromptTokensPath(String promptTokensPath) {
        this.promptTokensPath = promptTokensPath;
    }

    public String getCompletionTokensPath() {
        return completionTokensPath;
    }

    public void setCompletionTokensPath(String completionTokensPath) {
        this.completionTokensPath = completionTokensPath;
    }

    public JSONObjectParser<MessageStatus> getStatusParser() {
        return statusParser;
    }

    public void setStatusParser(JSONObjectParser<MessageStatus> statusParser) {
        this.statusParser = statusParser;
    }

    public JSONObjectParser<List<FunctionCall>> getCallsParser() {
        return callsParser;
    }

    public void setCallsParser(JSONObjectParser<List<FunctionCall>> callsParser) {
        this.callsParser = callsParser;
    }

    public List<FunctionCall> getCalls() {
        return calls;
    }

    public void setCalls(List<FunctionCall> calls) {
        this.calls = calls;
    }

    @Override
    public AiMessage parse(JSONObject rootJson) {
        AiMessage aiMessage = new AiMessage();

        if (StringUtil.hasText(this.contentPath)) {
            aiMessage.setContent((String) JSONPath.eval(rootJson, this.contentPath));
        }

        if (StringUtil.hasText(this.reasoningContentPath)) {
            aiMessage.setReasoningContent((String) JSONPath.eval(rootJson, this.reasoningContentPath));
        }

        if (StringUtil.hasText(this.indexPath)) {
            aiMessage.setIndex((Integer) JSONPath.eval(rootJson, this.indexPath));
        }


        if (StringUtil.hasText(promptTokensPath)) {
            aiMessage.setPromptTokens((Integer) JSONPath.eval(rootJson, this.promptTokensPath));
        }

        if (StringUtil.hasText(completionTokensPath)) {
            aiMessage.setCompletionTokens((Integer) JSONPath.eval(rootJson, this.completionTokensPath));
        }

        if (StringUtil.hasText(this.totalTokensPath)) {
            aiMessage.setTotalTokens((Integer) JSONPath.eval(rootJson, this.totalTokensPath));
        }
        //some LLMs like Ollama not response the total tokens
        else if (aiMessage.getPromptTokens() != null && aiMessage.getCompletionTokens() != null) {
            aiMessage.setTotalTokens(aiMessage.getPromptTokens() + aiMessage.getCompletionTokens());
        }

        if (this.statusParser != null) {
            aiMessage.setStatus(this.statusParser.parse(rootJson));
        }

        if (callsParser != null) {
            aiMessage.setCalls(callsParser.parse(rootJson));
        }

        return aiMessage;
    }


    public static DeepseekAiMessageParser getChatGPTMessageParser(boolean isStream) {
        DeepseekAiMessageParser aiMessageParser = new DeepseekAiMessageParser();
        if (isStream) {
            aiMessageParser.setContentPath("$.choices[0].delta.content");
            aiMessageParser.setReasoningContentPath("$.choices[0].delta.reasoning_content");
        } else {
            aiMessageParser.setContentPath("$.choices[0].message.content");
            aiMessageParser.setReasoningContentPath("$.choices[0].message.reasoning_content");
        }

        aiMessageParser.setIndexPath("$.choices[0].index");
        aiMessageParser.setTotalTokensPath("$.usage.total_tokens");
        aiMessageParser.setPromptTokensPath("$.usage.prompt_tokens");
        aiMessageParser.setCompletionTokensPath("$.usage.completion_tokens");

        aiMessageParser.setStatusParser(content -> {
            Object finishReason = JSONPath.eval(content, "$.choices[0].finish_reason");
            if (finishReason != null) {
                return MessageStatus.END;
            }
            return MessageStatus.MIDDLE;
        });

        aiMessageParser.setCallsParser(content -> {
            log.info(content.toString());

                return  getNoStream(content);

        });


        return aiMessageParser;
    }

    private static FunctionCall getFunctionCall( DeepseekAiMessageParser aiMessageParser,Integer index){
        FunctionCall functionCall = new FunctionCall();
        if(aiMessageParser.getCalls()!=null&&aiMessageParser.getCalls().size()>=index){
           return  aiMessageParser.getCalls().get(index);
        }
            return  new FunctionCall();

    }
    private static List<FunctionCall> getStream(DeepseekAiMessageParser aiMessageParser,JSONObject content){
        JSONArray toolCalls = (JSONArray) JSONPath.eval(content, "$.choices[0].message.tool_calls");
        if (toolCalls == null || toolCalls.isEmpty()) {
            return Collections.emptyList();
        }
        List<FunctionCall> functionCalls = new ArrayList<>();
        for (int i = 0; i < toolCalls.size(); i++) {
            JSONObject jsonObject = toolCalls.getJSONObject(i);
            JSONObject functionObject = jsonObject.getJSONObject("function");
            if (functionObject != null) {
                FunctionCall functionCall = aiMessageParser.getCalls().get(i);
                if(jsonObject.containsKey("id")) {
                    functionCall.setId(jsonObject.getString("id"));
                }
                if(functionObject.containsKey("name")) {
                    functionCall.setName(functionObject.getString("name"));
                }
                Object arguments = functionObject.get("arguments");
                if (arguments instanceof Map) {
                    //noinspection unchecked
                    functionCall.setArgs((Map<String, Object>) arguments);
                } else if (arguments instanceof String) {
                    //noinspection unchecked
                    functionCall.setArgs(JSON.parseObject(arguments.toString(), Map.class));
                }
                functionCalls.add(functionCall);
            }
        }
        return functionCalls;
    }

    private static List<FunctionCall> getNoStream(JSONObject content){
        JSONArray toolCalls = (JSONArray) JSONPath.eval(content, "$.choices[0].message.tool_calls");
        if (toolCalls == null || toolCalls.isEmpty()) {
            return Collections.emptyList();
        }
        List<FunctionCall> functionCalls = new ArrayList<>();
        for (int i = 0; i < toolCalls.size(); i++) {
            JSONObject jsonObject = toolCalls.getJSONObject(i);
            JSONObject functionObject = jsonObject.getJSONObject("function");
            if (functionObject != null) {
                FunctionCall functionCall = new FunctionCall();
                functionCall.setId(jsonObject.getString("id"));
                functionCall.setName(functionObject.getString("name"));
                Object arguments = functionObject.get("arguments");
                if (arguments instanceof Map) {
                    //noinspection unchecked
                   Map<String,Object> argu= functionCall.getArgs();
                    Map<String, Object> parame = (Map<String, Object>) arguments;
                   for (String key:parame.keySet()) {
                      if(argu.containsKey(key)){
                          argu.put(key,argu.get(key).toString()+parame.get(key).toString());
                      }else{
                          argu.put(key,parame.get(key).toString());
                      }
                   }
                    functionCall.setArgs(argu);
                } else if (arguments instanceof String) {
                    //noinspection unchecked
                    Map<String,Object> argu= functionCall.getArgs();
                    if(argu==null){
                        argu=new HashMap<>();
                    }
                    argu.put("arguments",arguments);
                    functionCall.setArgs(JSON.parseObject(arguments.toString(), Map.class));
                }
                functionCalls.add(functionCall);
            }
        }
        return functionCalls;
    }

    public static void main(String[] argc){
        String content="{\"id\":\"34b011dc-06d0-4c68-9114-61f6efa3eae1\",\"object\":\"chat.completion.chunk\",\"created\":1744624255,\"model\":\"deepseek-chat\",\"system_fingerprint\":\"fp_3d5141a69a_prod0225\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"获取\"}}]},\"logprobs\":null,\"finish_reason\":null}]}";

        JSONArray toolCalls = (JSONArray) JSONPath.eval(content, "$.choices[0].delta.tool_calls");
        System.out.println(toolCalls);
        for (int i = 0; i < toolCalls.size(); i++) {
            JSONObject jsonObject = toolCalls.getJSONObject(i);
            System.out.println(jsonObject);
            JSONObject functionObject = jsonObject.getJSONObject("function");
            Object arguments = functionObject.get("arguments");
            System.out.println(arguments);
        }
    }
}
